from flask import Flask, request, send_file, make_response, jsonify, render_template, render_template_string, session  
from your_image_generation_library import generate_image_from_parameters
from PIL import Image
import io  
import base64
from waitress import serve   
from flask_wtf import FlaskForm  
from wtforms import StringField, PasswordField, SubmitField  
from wtforms.validators import DataRequired, Length, EqualTo  
from captcha.image import ImageCaptcha  
import random  
import secrets  

app = Flask(__name__)  

# 假设的用户数据库（实际应用中应使用数据库）  
users = {  
    'user1': 'password1',  
    'user2': 'password2'  
}  


secret_key = secrets.token_hex(16)
app.secret_key = secret_key  # 用于会话  

# 验证码生成器  
captcha_generator = ImageCaptcha(width=200, height=80)  

# 存储验证码的session key  
CAPTCHA_KEY = 'captcha' 

@app.route('/')  
def home():  
    return render_template('index.html')  # 返回HTML模板


@app.route('/login', methods=['POST'])  
def login():  
    # 确保请求是POST并且包含JSON数据  
    if request.method != 'POST' or not request.is_json:  
        return jsonify({"error": "Invalid request"}), 400  
  
    # 从JSON中获取用户信息  
    login_info = request.get_json()  
    username = login_info.get('username')  
    password = login_info.get('password')  
  
    # 验证用户名和密码  
    if username in users and users[username] == password:  
        # 登录成功，返回token或其他成功信息（这里仅返回成功信息）  
        return jsonify({"message": "Login successful"}), 200  
    else:  
        # 登录失败，返回错误信息  
        return jsonify({"error": "Invalid credentials"}), 401  


class RegistrationForm(FlaskForm):  
    username = StringField('Username', validators=[DataRequired(), Length(min=4, max=25)])  
    password = PasswordField('Password', validators=[DataRequired(), Length(min=6)])  
    confirm_password = PasswordField('Confirm Password', validators=[DataRequired(), EqualTo('password')])  
    captcha = StringField('Captcha', validators=[DataRequired()])  
  
@app.route('/register', methods=['GET', 'POST'])  
def register():  
    form = RegistrationForm()  
    captcha_text = ''.join(random.choices('ABCDEFGHJKLMNPQRSTUVWXYZ23456789', k=6))  
    session[CAPTCHA_KEY] = captcha_text  
    captcha_image = captcha_generator.generate_image(captcha_text)  
    captcha_url = 'data:image/png;base64,' + captcha_image.encode('base64').decode()  
  
    if request.method == 'GET':  
        return jsonify({"captcha_url": captcha_url})
  
    if form.validate_on_submit():  
        # 表单验证成功，这里可以添加将用户保存到数据库的代码  
        # 检查验证码是否正确  
        if form.captcha.data.lower() == session[CAPTCHA_KEY].lower():  
            # 验证码正确，处理注册逻辑（如保存到数据库）  
            return jsonify({"message": "Registration successful"}), 200  
        else:  
            return jsonify({"error": "Invalid captcha"}), 400  
  
    # 表单验证失败或验证码错误  
    return jsonify({"error": "Form validation failed"}), 400  



@app.route('/api/images/generate', methods=['POST'])
def generate_image():  
    data = request.get_json()
    prompt = data.pop('prompt','best quality')  
    # image = data.get('image', None)  
    width = int(data.pop('width', None))  
    height = int(data.pop('height', None))  
    num_inference_steps = int(data.get('num_inference_steps', 30))
    guidance_scale = float(data.get('guidance_scale', 3.5))
  
    # 生成图片
    image = generate_image_from_parameters(prompt, width, height, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps)  

    # 检查图像是否为None或不是PIL.Image对象  
    if image is None or not isinstance(image, Image.Image):  
        return jsonify({"error": "Unable to generate image"}), 400 
  
    # 将 PIL Image 对象转换为字节流  
    img_byte_arr = io.BytesIO()  
    image.save(img_byte_arr, format='JPEG')  
    img_byte_arr = img_byte_arr.getvalue()  
  
    # 设置响应类型为图像  
    base64_img_str = base64.b64encode(img_byte_arr).decode('utf-8')  
  
    # 返回Base64编码的图像字符串  
    return jsonify({"image": f"data:image/jpeg;base64,{base64_img_str}"})   


'''
# 调试用
if __name__ == '__main__':  
    app.run()
'''

# 部署服务
if __name__ == '__main__':  
    serve(app.wsgi_app, host='0.0.0.0', port=8000)

